#include "txr_algo_dlm_cls_v1.h"
#include "yolov6.h"
#include "common\define.h"

namespace txr_algo_dlm_cls
{
	TxrAlgoDlmClsV1::TxrAlgoDlmClsV1()
		:m_p_detect(TK_NULL_PTR)
	{

	}

	TxrAlgoDlmClsV1::~TxrAlgoDlmClsV1()
	{
		Release();
	}

	tk_bool TxrAlgoDlmClsV1::Load(st_dlm_cfg & cfg)
	{
		m_p_detect = new DlmClsV1();
		std::string path_weight = std::string(cfg.file_path_name) + SUFFIX_V1;
		return m_p_detect->Init(cfg.gpu_id, path_weight.c_str());
	}

	void TxrAlgoDlmClsV1::Release()
	{
		SAFE_DELETE(m_p_detect);
	}

	tk_int32 TxrAlgoDlmClsV1::BatchSize()
	{
		if (m_p_detect)
		{
			return m_p_detect->BatchSize();
		}
		return 0;
	}

	void TxrAlgoDlmClsV1::Detect(st_detect_unit * p_unit, tk_int32 size)
	{
		if (!size)
		{
			return;
		}
		for (tk_int32 i = 0; i < size; ++i)
		{
			st_img_rgb & rgb = p_unit[i].img;
			if (!
				((rgb.pp_rgb8[0] && rgb.pp_rgb8[1] && rgb.pp_rgb8[2]) ||
				(rgb.pp_rgb16[0] && rgb.pp_rgb16[1] && rgb.pp_rgb16[2]) ||
					(rgb.pp_rgbf[0] && rgb.pp_rgbf[1] && rgb.pp_rgbf[2]))
				)
			{
				Q_ASSERT(0);
				return;
			}
		}

		std::vector<st_dlm_data> v_dlm_data;
		v_dlm_data.resize(size);

		for (tk_int32 i = 0; i < size; ++i)
		{
			st_img_rgb & rgb = p_unit[i].img;
			image_t & img = v_dlm_data[i].img;
			if (rgb.pp_rgbf[0])
			{
				img.rgbf = rgb.pp_rgbf;
			}
			else if (rgb.pp_rgb8[0])
			{
				img.rgb8 = rgb.pp_rgb8;
			}
			else if (rgb.pp_rgb16[0])
			{
				img.rgb16 = rgb.pp_rgb16;
			}

			img.h = rgb.h;
			img.w = rgb.w;
			img.c = rgb.c;

			v_dlm_data[i].cls_prob_thres = p_unit[i].param.prob_thres;
		}
		std::vector<st_result_info> v_results;
		m_p_detect->Detect(v_dlm_data.data(), size, v_results);

		for (int i = 0; i < size; ++i)
		{
			p_unit[i].result = v_results[i];
		}
	}
	//===========================================// extern "C" 表示使用C语言链接约定，这样其他语言可以调用这些函数
	// TXR_ALGO_DLM_CLS_V1_EXPORT 是导出宏，用于将函数导出到DLL
	extern "C" TXR_ALGO_DLM_CLS_V1_EXPORT TxrAlgoDlmClsImpl* CreateDlmClsObj();
	extern "C" TXR_ALGO_DLM_CLS_V1_EXPORT void DeleteDlmClsObj(TxrAlgoDlmClsImpl* p);
	TxrAlgoDlmClsImpl* CreateDlmClsObj()
	{
		return new TxrAlgoDlmClsV1();
	}

	void DeleteDlmClsObj(TxrAlgoDlmClsImpl* p)
	{
		SAFE_DELETE(p);
	}
}
